package com.example.ai.service;

import com.example.ai.entity.Document;
import com.example.ai.repository.DocumentRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import java.io.IOException;
import java.util.List;

@Service
@RequiredArgsConstructor
public class DocumentService {
    private final DocumentRepository documentRepository;
    private final BaiduOcrService ocrService;

    public Document uploadDocument(MultipartFile file, Long userId) {
        try {
            String content;
            if (isPdf(file)) {
                content = extractPdfText(file);
            } else {
                content = ocrService.recognizeText(file);
            }

            Document document = new Document();
            document.setFileName(file.getOriginalFilename());
            document.setContent(content);
            document.setUserId(userId);

            return documentRepository.save(document);
        } catch (IOException e) {
            throw new RuntimeException("文件处理失败", e);
        }
    }

    public List<Document> getUserDocuments(Long userId) {
        return documentRepository.findByUserId(userId);
    }

    public Document getDocument(Long id) {
        return documentRepository.findById(id)
            .orElseThrow(() -> new RuntimeException("文档不存在"));
    }

    private boolean isPdf(MultipartFile file) {
        return file.getContentType() != null && 
               file.getContentType().equals("application/pdf");
    }

    private String extractPdfText(MultipartFile file) throws IOException {
        try (PDDocument document = PDDocument.load(file.getInputStream())) {
            PDFTextStripper stripper = new PDFTextStripper();
            return stripper.getText(document);
        }
    }
} 